package analysis;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;

import analysis.NFAAnalyser.IdaSpecialTransitionLabel;
import analysis.NFAAnalyserInterface.EdaAnalysisResultsESCC;
import analysis.NFAAnalyserInterface.EdaAnalysisResultsFilter;
import analysis.NFAAnalyserInterface.EdaAnalysisResultsParallel;
import analysis.NFAAnalyserInterface.IdaAnalysisResultsIda;

import nfa.NFAGraph;
import nfa.NFAEdge;
import nfa.NFAVertexND;
import nfa.transitionlabel.*;
import nfa.transitionlabel.TransitionLabel.TransitionType;

public class ExploitStringBuilder implements ExploitStringBuilderInterface<EdaAnalysisResults, IdaAnalysisResults> {

	@Override
	public ExploitString buildEdaExploitString(EdaAnalysisResults results) throws InterruptedException {
		
		switch (results.edaCase) {
		case PARALLEL:
			return getParallelExploitString((EdaAnalysisResultsParallel) results);
		case ESCC:
			return getEsccExploitString((EdaAnalysisResultsESCC) results);
		case FILTER:
			return getFilterExploitString((EdaAnalysisResultsFilter) results);
		case NO_EDA:
			throw new NoExploitStringException("The graph does not have EDA.");
		default:
			throw new RuntimeException("Invalid case for EDA");
		}
	}
	
	@Override
	public ExploitString buildIdaExploitString(IdaAnalysisResults results) throws InterruptedException {
		
		switch (results.idaCase) {
		case IDA:
			return getIdaExploitString((IdaAnalysisResultsIda) results);
		case NO_IDA:
			throw new NoExploitStringException("The graph does not have EDA.");
		}
		
		return null;
		
	}
	
	public static ExploitString getParallelExploitString(EdaAnalysisResultsParallel parallelResults) throws InterruptedException {
		NFAGraph originalGraph = parallelResults.getOriginalGraph();
		NFAVertexND sourceVertex = parallelResults.getSourceVertex();
		NFAGraph mergedScc = parallelResults.getMergedScc();
		NFAEdge parallelEdge = parallelResults.getParallelEdge();

		String prefixString = buildPrefixString(originalGraph, sourceVertex);
		NFAEdge incomingEdge = mergedScc.incomingEdgesOf(sourceVertex).iterator().next();
		String pumpString = buildPumpMultiPath(mergedScc, incomingEdge, parallelEdge);
		String suffixString = buildSuffixString(originalGraph, parallelEdge.getTargetVertex());
		return new ExploitString(prefixString, pumpString, suffixString);
	}
	
	public static ExploitString getEsccExploitString(EdaAnalysisResultsESCC esccResults) throws InterruptedException {
		NFAGraph originalGraph = esccResults.getOriginalGraph();
		NFAGraph originalScc = esccResults.getOriginalScc();
		NFAEdge entranceEdge = esccResults.getEntranceEdge();
		NFAEdge exitEdge = esccResults.getExitEdge();
		
		/* building the exploit string */
		NFAVertexND startVertex = entranceEdge.getTargetVertex();
		String prefixString = buildPrefixString(originalGraph, startVertex);

		String pumpString = buildPumpMultiPath(originalScc, entranceEdge, exitEdge);
		String suffixString = buildSuffixString(originalGraph, exitEdge.getTargetVertex());
		return new ExploitString(prefixString, pumpString, suffixString);
	}
	
	public static ExploitString getFilterExploitString(EdaAnalysisResultsFilter filterResults) throws InterruptedException {
		NFAGraph originalGraph = filterResults.getOriginalGraph();
		NFAVertexND endState = filterResults.getEndState();
		NFAVertexND startState = filterResults.getStartState();
		NFAGraph pcScc = filterResults.getPcScc();
		/* building the exploit string */
		String prefixString = buildPrefixString(originalGraph, startState.getStateByDimension(1));		
		String pumpString = buildPumpIntersect(pcScc, startState, endState);
		//System.out.println("IN ExploitStringBuilder:getFilterExploitString()");
		//System.out.println(originalGraph);
		//System.out.println("End");
		String suffixString = buildSuffixString(originalGraph, startState.getStateByDimension(1));
		return new ExploitString(prefixString, pumpString, suffixString);
	}
	
	/**
	 * Builds the string that takes the NFA to the state that shows EDA.
	 * 
	 * @param m
	 *            The graph reperesenting the NFA.
	 * @param finish
	 *            The state that shows EDA.
	 * @return The string built.
	 */
	public static String buildPrefixString(NFAGraph m, NFAVertexND finish) {
		LinkedList<NFAEdge> edges = NFAAnalysisTools.shortestPathTo(m, finish);

		return buildStringFromEdges(edges);
	}

	public static String buildSuffixString(NFAGraph n, NFAVertexND start) throws InterruptedException {
		/* determining the alphabet */
		//System.out.println("ExploitStringBuilder:Start");
		HashSet<TransitionLabel> regexAlphabet = (HashSet<TransitionLabel>) NFAAnalysisTools.getAlphabet(n);
		//System.out.println("ExploitStringBuilder:1");
		/* constructing n' */
		NFAGraph nAccent = n;
		

		//System.out.println("IN ExploitStringBuilder:buildSuffixString()");
		//System.out.println(n);
		//System.out.println("End");
		/* determinizing */
		NFAGraph dfa = NFAAnalysisTools.determinize(nAccent, n.vertexSet(), regexAlphabet);	
		/* swapping final and nonfinal states */
		//System.out.println("ExploitStringBuilder:2");
		
		for (NFAVertexND currentState : dfa.vertexSet()) {
			if (dfa.isAcceptingState(currentState)) {
				dfa.removeAcceptingState(currentState);
			} else {
				dfa.addAcceptingState(currentState);
			}
		}
		//System.out.println("ExploitStringBuilder:3");
		
		if (dfa.getAcceptingStates().isEmpty()) {
			
			
			Iterator<TransitionLabel> i0 = regexAlphabet.iterator();
			TransitionLabel wholeAlphabet = i0.next();
			while (i0.hasNext()) {
				wholeAlphabet = wholeAlphabet.union(i0.next());
			}
			
			// if the alphabet matches our default character, find another character 
			if (!wholeAlphabet.matches("#")) {
				return "#";  
			} else {	
				
				/* If there is any character in the 16UNICODE alphabet that isn't in the regex alpahbet, take this character */
				HashSet<TransitionLabel> unicode16Alphabet = new HashSet<TransitionLabel>();
				for (int i = 0; i < TransitionLabel.MAX_16UNICODE; i++) {
					char currentChar = (char) i;
					String currentString = "" + currentChar;
					if (currentChar == '[') {
						currentString = "\\" + currentChar;
					} else if (currentChar == '\u03b5') {
						currentString = "[\u03b5]";
					}
					if (!wholeAlphabet.matches(currentString)) {						
						return "" + ((char) i);
					}
				}

				/* Otherwise, determinize from the state where EDA/IDA took place */
				unicode16Alphabet.add(CharacterClassTransitionLabel.wildcardLabel());
				// XXX I'm not sure if this is correct.
				// We determinize starting at the state where EDA/IDA occurs.
				// Then we take the complement of this DFA
				// and try to find a nonempty suffix, 
				// since for an empty suffix, the matcher might spit out part of the pump and use it to accept the string: (a.*)|(((ba)|.??)*a)
				HashSet<NFAVertexND> startStates = NFAAnalysisTools.reachableWithEpsilon(nAccent, start); 

				NFAGraph dfa2 = NFAAnalysisTools.determinize(nAccent, startStates, unicode16Alphabet);
				/* swapping final and nonfinal states */;
				for (NFAVertexND currentState : dfa2.vertexSet()) {
					if (dfa2.isAcceptingState(currentState)) {
						dfa2.removeAcceptingState(currentState);
					} else {
						dfa2.addAcceptingState(currentState);
						//System.out.println("Check: " + currentState);
					}
				}

				if (!dfa2.getAcceptingStates().isEmpty()) {
					Iterator<NFAVertexND> i1 = dfa2.getAcceptingStates().iterator();
					while (i1.hasNext()) {
					NFAVertexND finalState = i1.next();
						//System.out.println(finalState);	
		
						LinkedList<NFAEdge> pathToFinal = NFAAnalysisTools.shortestPathTo(dfa2, finalState);
						//for (NFAEdge e : pathToFinal) {
						//	System.out.print(e.getSourceVertex() + "-" + e + "->" + e.getTargetVertex());
						//}
						//System.out.println();

						if (!pathToFinal.isEmpty()) {
							return buildStringFromEdges(pathToFinal);
						}
					}
		

					throw new RuntimeException("Cannot build suffix! (Only ε suffixes.)");
				} else {
					throw new RuntimeException("Cannot build suffix!");
				}
				
				
			}
			//return "#";
		}
		
		
		Iterator<NFAVertexND> i0 = dfa.getAcceptingStates().iterator();
		NFAVertexND finalState = i0.next();
		
		LinkedList<NFAEdge> pathToFinal = NFAAnalysisTools.shortestPathTo(dfa, finalState);
		
		//System.out.println("ExploitStringBuilder:End");
		return buildStringFromEdges(pathToFinal);
	}

	/**
	 * This function loops through a list of edges and build a string of all
	 * edges representing a symbol transition.
	 * 
	 * @param edges
	 *            The list of edges.
	 * @return The string built.
	 */
	public static String buildStringFromEdges(LinkedList<NFAEdge> edges) {

		if (edges == null) {
			return null;
		}

		StringBuilder toReturn = new StringBuilder();

		for (NFAEdge e : edges) {
			if (!e.getIsEpsilonTransition()) {
				toReturn.append(e.getATransitionCharacter());
			}
		}

		return toReturn.toString();
	}

	private static String buildPumpMultiPath(NFAGraph scc, NFAEdge entranceEdge, NFAEdge exitEdge) {
		//System.out.println(scc);
		String toReturn;
		LinkedList<LinkedList<NFAEdge>> pathsToP = new LinkedList<LinkedList<NFAEdge>>();
		LinkedList<NFAEdge> queue = new LinkedList<NFAEdge>();
		HashMap<NFAEdge, LinkedList<NFAEdge>> paths = new HashMap<NFAEdge, LinkedList<NFAEdge>>();
		HashSet<NFAEdge> visitedEdges = new HashSet<NFAEdge>();

		queue.add(exitEdge); /* start at the exit edge */
		LinkedList<NFAEdge> newPath = new LinkedList<NFAEdge>();
		newPath.add(exitEdge);
		paths.put(exitEdge, newPath);
		visitedEdges.add(exitEdge);
		while (!queue.isEmpty()) {

			NFAEdge currentEdge = queue.removeLast();
			LinkedList<NFAEdge> currentPath = paths.get(currentEdge);
			if (entranceEdge.equals(currentEdge)) {
				newPath = new LinkedList<NFAEdge>(currentPath);
				newPath.add(currentEdge);
				pathsToP.add(newPath);
				break;
			}
			NFAVertexND targetVertex = currentEdge.getTargetVertex();
			//if (targetVertex.equals(new NFAVertexND("q35"))) {
			//	for (NFAEdge e : scc.outgoingEdgesOf(targetVertex)) {
			//		System.out.println("Outgoing edge: " + e.getSourceVertex() + "-" + e + "->" + e.getTargetVertex());
			//	}
			//}

			for (NFAEdge e : scc.outgoingEdgesOf(targetVertex)) {
				if (!visitedEdges.contains(e)) {
					visitedEdges.add(e);
					newPath = new LinkedList<NFAEdge>(currentPath);
					newPath.add(e);
					paths.put(e, newPath);
					queue.addFirst(e);
				}

			}

		}
		//System.out.println("Entrance: " + entranceEdge.getSourceVertex() + "-" + entranceEdge + "->" + entranceEdge.getTargetVertex());
		//System.out.println("Exit: " + exitEdge.getSourceVertex() + "-" + exitEdge + "->" + exitEdge.getTargetVertex());
		//for (LinkedList<NFAEdge> pathToP : pathsToP) {
		//	System.out.println();
		//	for (NFAEdge e : pathToP) {
		//		System.out.print(e.getSourceVertex() + "-" + e + "->" + e.getTargetVertex() + " ; ");
		//	}
		//	System.out.println();
		//}
		toReturn = buildStringFromEdges(pathsToP.get(0));

		return toReturn;
	}

	/**
	 * This function builds the exploit string for a vulnerable regular
	 * expression, in the cases where p and q were found in the scc.
	 * 
	 * @param pcscc
	 *            The scc in the product construction.
	 * @param p
	 *            The state that shows EDA, to build the path from.
	 * @param q
	 *            The state to build the path to.
	 * @return The string capable of exploiting the vulnerable regular
	 *         expression.
	 */
	public static String buildPumpIntersect(NFAGraph pcscc, NFAVertexND p, NFAVertexND q) {
		
		//LinkedList<NFAEdge> pathToQ = buildExploitStringMultiPathTo(pcscc, pcscc.incomingEdgesOf(q).iterator().next(), pcscc.outgoingEdgesOf(p).iterator().next());
		LinkedList<NFAEdge> pathToQ = buildExploitStringSCCPathTo(pcscc, p, q);
		String s1 = buildStringFromEdges(pathToQ);
		StringBuilder toReturn = new StringBuilder(s1);
		//LinkedList<NFAEdge> pathToP = buildExploitStringMultiPathTo(pcscc, pcscc.incomingEdgesOf(p).iterator().next(), pcscc.outgoingEdgesOf(q).iterator().next());
		LinkedList<NFAEdge> pathToP = buildExploitStringSCCPathTo(pcscc, q, p);
		
		String s2 = buildStringFromEdges(pathToP);
		toReturn.append(s2);
		return toReturn.toString();
	}

	/**
	 * This function finds a path from a state to another.
	 * 
	 * @param scc
	 *            The strongly connected component to find the path in.
	 * @param p
	 *            The state to search from.
	 * @param q
	 *            The state to search to.
	 * @return A linked list containing the edges in the path.
	 */
	private static LinkedList<NFAEdge> buildExploitStringSCCPathTo(NFAGraph scc, NFAVertexND p, NFAVertexND q) {
		LinkedList<NFAVertexND> queue = new LinkedList<NFAVertexND>();
		HashMap<NFAVertexND, LinkedList<NFAEdge>> paths = new HashMap<NFAVertexND, LinkedList<NFAEdge>>();
		HashSet<NFAEdge> visitedEdges = new HashSet<NFAEdge>();

		for (NFAEdge e : scc.outgoingEdgesOf(p)) {
			NFAVertexND target = e.getTargetVertex();

			queue.addFirst(target);
			visitedEdges.add(e);
			LinkedList<NFAEdge> newPath = new LinkedList<NFAEdge>();
			newPath.add(e);
			paths.put(target, newPath);
			if (target.equals(q)) {
				return newPath;
			}
		}

		while (!queue.isEmpty()) {
			NFAVertexND current = queue.removeLast();
			LinkedList<NFAEdge> currentPath = paths.get(current);

			for (NFAEdge e : scc.outgoingEdgesOf(current)) {
				if (!visitedEdges.contains(e)) {
					visitedEdges.add(e);
					NFAVertexND target = e.getTargetVertex();
					LinkedList<NFAEdge> newPath = new LinkedList<NFAEdge>(currentPath);
					newPath.add(e);
					paths.put(target, newPath);
					queue.addFirst(target);
					if (target.equals(q)) {
						return newPath;
					}
				}

			}

		}
		return null;
	}

	private static ExploitString getIdaExploitString(IdaAnalysisResultsIda idaResults) throws InterruptedException {
		NFAGraph originalGraph = idaResults.originalGraph;
		int degree = idaResults.getDegree();
		String[] pumps = new String[degree];
		String[] separators = new String[degree];
		int degreeCounter = 0;
		LinkedList<NFAEdge> maxPath = idaResults.getMaxPath();
		Iterator<NFAEdge> i0 = maxPath.iterator();
		
		while (i0.hasNext()) {	
			TransitionLabel currentTransitionLabel = i0.next().getTransitionLabel();
			StringBuilder separatorBuilder = new StringBuilder("");
			while (!(currentTransitionLabel instanceof IdaSpecialTransitionLabel)) {
				if (currentTransitionLabel.getTransitionType() != TransitionType.EPSILON) {
					separatorBuilder.append(currentTransitionLabel.getSymbol());
				}
				
				if (i0.hasNext()) {	
					currentTransitionLabel = i0.next().getTransitionLabel();
				} else {
					break;
				}
				
			}
			/* If nothing follows, the last "separator" was actually a valid suffix */
			if (i0.hasNext()) {
				
				separators[degreeCounter] = separatorBuilder.toString();
				
				/* One IdaSpecialTransitionLabel contains a pump, we assume it may not contain an epsilon */
				//pumps[degreeCounter] = currentTransitionLabel.getSymbol();
				StringBuilder pumpBuilder = new StringBuilder();
				IdaSpecialTransitionLabel currentIdaSpecialTransitionLabel = (IdaSpecialTransitionLabel) currentTransitionLabel;
				for (TransitionLabel transitionLabelInSpecialTransitionLabel : currentIdaSpecialTransitionLabel.getTransitionLabels()) {
					if (transitionLabelInSpecialTransitionLabel.getTransitionType() != TransitionType.EPSILON) {
						pumpBuilder.append(transitionLabelInSpecialTransitionLabel.getSymbol());
					}
				}	
				pumps[degreeCounter] = pumpBuilder.toString();
				degreeCounter++;
			}				
		}
		
		NFAVertexND suffixStartVertex = maxPath.get(maxPath.size() - 1).getTargetVertex();
		//System.out.println("check1");
		String suffix = buildSuffixString(originalGraph, suffixStartVertex);
		//System.out.println("check2");
		ExploitString es = new ExploitString(separators, pumps, suffix);
		return es;
		
	}
	
}
